# apps/orchestrator/async_orchestrator_agent.py
import asyncio
import aiohttp
from typing import Optional, Dict, Any, List, Tuple
import asyncio
import aiohttp
from datetime import datetime
import json
import threading
from loguru import logger

# 导入项目核心模块和常量
from utils.logger import Logger
from utils.exceptions import (
    RedisError, LockAcquireError, ParamError,
    AgentNotFoundError, AgentCallError
)
from utils.constants import (
    AgentType, TaskState, ErrorCode,
    AGENT_STATE_UPDATE_CHANNEL, TASK_STATE_UPDATE_CHANNEL
)
from middleware.redis_adapter import get_redis_adapter
from states.agent_state import AgentState
from states.task_state import TaskStateManager
from apps.orchestrator.task_reviewer import task_reviewer


class AsyncOrchestratorAgent:
    """异步编排Agent（多Agent协作核心调度器）"""

    def __init__(self, redis_client=None):  # 兼容原有接口，现在使用RedisAdapter
        # 依赖组件
        self.redis_client = redis_client or get_redis_adapter()
        self.logger = Logger.get_logger(agent_name="async_orchestrator_agent")

        # 缓存：健康Agent实例（减少Redis查询）
        self.healthy_agents: Dict[str, List[Dict[str, Any]]] = {}  # {agent_type: [instance_info]}
        self.cache_lock = threading.Lock()  # 缓存更新锁（线程安全）

        # 异步HTTP会话（连接池复用）
        self._session: Optional[aiohttp.ClientSession] = None
        
        # 启动Agent状态订阅（实时感知实例上下线）
        self._start_agent_state_subscribe()
        self.logger.success("异步编排Agent初始化完成，已启动Agent状态订阅")

    async def __aenter__(self):
        """异步上下文管理器进入"""
        if not self._session or self._session.closed:
            self._session = aiohttp.ClientSession(
                timeout=aiohttp.ClientTimeout(total=30),
                connector=aiohttp.TCPConnector(
                    limit=100,  # 连接池大小
                    limit_per_host=20,  # 每个主机最大连接数
                    ttl_dns_cache=300,  # DNS缓存时间
                    use_dns_cache=True,
                ),
                headers={
                    'User-Agent': 'MAESS-AsyncOrchestrator/1.0',
                    'Content-Type': 'application/json',
                }
            )
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """异步上下文管理器退出"""
        if self._session and not self._session.closed:
            await self._session.close()

    @property
    async def session(self) -> aiohttp.ClientSession:
        """获取或创建HTTP会话"""
        if not self._session or self._session.closed:
            self._session = aiohttp.ClientSession(
                timeout=aiohttp.ClientTimeout(total=30),
                connector=aiohttp.TCPConnector(
                    limit=100,
                    limit_per_host=20,
                    ttl_dns_cache=300,
                    use_dns_cache=True,
                ),
                headers={
                    'User-Agent': 'MAESS-AsyncOrchestrator/1.0',
                    'Content-Type': 'application/json',
                }
            )
        return self._session
        
    async def _get_healthy_agents(self) -> Dict[str, Dict[str, Any]]:
        """
        获取所有健康的Agent实例
        返回格式: {agent_id: {agent_type, capabilities, task_count, ...}}
        """
        result = {}
        with self.cache_lock:
            for agent_type, instances in self.healthy_agents.items():
                for instance in instances:
                    agent_id = instance["instance_id"]
                    result[agent_id] = {
                        "agent_type": agent_type,
                        "capabilities": [],  # 可以从Redis获取实际能力
                        "task_count": instance.get("current_tasks", 0),
                        "host_ip": instance["host_ip"],
                        "port": instance["port"]
                    }
        return result

    def _start_agent_state_subscribe(self) -> None:
        """启动Agent状态订阅（实时更新健康实例缓存）"""
        try:
            # 创建停止事件
            self._stop_event = threading.Event()

            def agent_state_callback(data: Dict[str, Any]) -> None:
                """Agent状态更新回调（实例注册/注销/更新时触发）"""
                try:
                    agent_type = data.get("agent_type", "")
                    instance_id = data.get("instance_id", "")
                    action = data.get("action", "update")
                    status = data.get("status", "")

                    self.logger.debug(
                        f"收到Agent状态更新: action={action} | type={agent_type} | instance={instance_id} | status={status}")

                    # 加锁更新健康实例缓存
                    with self.cache_lock:
                        if action in ["register", "update"] and status == "healthy":
                            # 新增/更新健康实例
                            if agent_type not in self.healthy_agents:
                                self.healthy_agents[agent_type] = []
                            # 移除旧实例信息（避免重复）
                            self.healthy_agents[agent_type] = [
                                ins for ins in self.healthy_agents[agent_type]
                                if ins["instance_id"] != instance_id
                            ]
                            # 添加新实例信息
                            self.healthy_agents[agent_type].append({
                                "instance_id": instance_id,
                                "host_ip": data.get("host_ip", ""),
                                "port": data.get("port", 0),
                                "last_heartbeat": data.get("last_heartbeat", 0)
                            })
                        elif action == "unregister" or status != "healthy":
                            # 移除注销/不健康实例
                            if agent_type in self.healthy_agents:
                                self.healthy_agents[agent_type] = [
                                    ins for ins in self.healthy_agents[agent_type]
                                    if ins["instance_id"] != instance_id
                                ]
                                # 空列表时删除键
                                if not self.healthy_agents[agent_type]:
                                    del self.healthy_agents[agent_type]
                except Exception as e:
                    self.logger.error(f"处理Agent状态更新失败: {str(e)}")

            # 启动订阅线程（确保使用正确的subscribe方法参数：频道、回调函数、停止事件）
            subscribe_thread = threading.Thread(
                target=self.redis_client.subscribe,
                args=(AGENT_STATE_UPDATE_CHANNEL, agent_state_callback, self._stop_event),
                daemon=True
            )
            subscribe_thread.start()
            self.logger.info(f"Agent状态订阅线程已启动，频道={AGENT_STATE_UPDATE_CHANNEL}")
            
        except Exception as e:
            self.logger.error(f"启动Agent状态订阅失败: {str(e)}")
            # 如果发生错误，尝试重新订阅
            import time
            if not hasattr(self, '_stop_event') or not self._stop_event.is_set():
                time.sleep(5)  # 等待5秒后重试
                self._start_agent_state_subscribe()

    def _select_best_agent_instance(self, agent_type: str) -> Optional[Dict[str, Any]]:
        """选择最优Agent实例（负载均衡+健康状态优先）"""
        # 1. 先从缓存获取健康实例，缓存为空时从Redis刷新
        with self.cache_lock:
            healthy_instances = self.healthy_agents.get(agent_type, [])

        if not healthy_instances:
            self.logger.warning(f"缓存中无{agent_type}健康实例，从Redis刷新")
            healthy_instances = AgentState.get_healthy_instances(agent_type, self.redis_client)
            # 更新缓存
            with self.cache_lock:
                self.healthy_agents[agent_type] = healthy_instances

        # 如果仍然没有健康实例，抛出异常
        if not healthy_instances:
            self.logger.error(f"没有找到健康的{agent_type}实例")
            raise AgentNotFoundError(f"No healthy {agent_type} agent instances available")

        # 2. 负载均衡策略：选择当前任务数最少的实例
        # （从Redis获取最新的current_tasks，确保负载准确）
        best_instance = None
        min_task_count = float("inf")
        for instance in healthy_instances:
            instance_id = instance["instance_id"]
            redis_key = f"agent:state:{agent_type}:{instance_id}"
            # 从Redis获取当前任务数
            current_tasks = self.redis_client.hget(redis_key, "current_tasks")
            current_tasks = int(current_tasks) if current_tasks else 0

            # 选择任务数最少的实例
            if current_tasks < min_task_count:
                min_task_count = current_tasks
                best_instance = instance
                best_instance["current_tasks"] = current_tasks

        self.logger.debug(f"为{agent_type}选择最优实例: {best_instance['instance_id']} (当前任务数: {min_task_count})")
        return best_instance

    async def _call_agent_instance_async(self, instance: Dict[str, Any], task_params: Dict[str, Any]) -> Dict[str, Any]:
        """异步调用Agent实例执行任务（使用aiohttp实现真正的异步HTTP调用）"""
        agent_type = task_params.get("agent_type", "")
        instance_id = instance["instance_id"]
        
        host_ip = instance["host_ip"]
        port = instance["port"]

        self.logger.info(f"异步调用Agent实例: {instance_id} | 参数: {task_params}")

        # 构建异步HTTP请求URL和参数
        url = f"http://{host_ip}:{port}/api/run"
        request_data = {
            "task_id": task_params.get("task_id", ""),
            "inputs": {
                "query": task_params.get("input", ""),
                "session_id": task_params.get("session_id", f"session_{task_params.get('task_id', '')}")
            },
            "agent_params": {}
        }

        try:
            session = await self.session
            async with session.post(url, json=request_data) as response:
                # 设置超时时间
                response_data = await asyncio.wait_for(response.json(), timeout=30.0)
                
                if response.status == 200:
                    return {
                        "success": True,
                        "data": response_data,
                        "error_code": 0,
                        "error_msg": ""
                    }
                else:
                    return {
                        "success": False,
                        "data": None,
                        "error_code": response.status,
                        "error_msg": f"HTTP {response.status}: {response_data.get('error', 'Unknown error')}"
                    }

        except asyncio.TimeoutError:
            self.logger.error(f"调用Agent实例{instance_id}超时")
            return {
                "success": False,
                "data": None,
                "error_code": ErrorCode.AGENT_CALL_ERROR.value,
                "error_msg": "Agent调用超时(30s)"
            }
        except aiohttp.ClientError as e:
            self.logger.error(f"调用Agent实例{instance_id}网络错误: {str(e)}")
            return {
                "success": False,
                "data": None,
                "error_code": ErrorCode.AGENT_CALL_ERROR.value,
                "error_msg": f"网络错误: {str(e)}"
            }
        except Exception as e:
            self.logger.error(f"调用Agent实例{instance_id}失败: {str(e)}")
            return {
                "success": False,
                "data": None,
                "error_code": ErrorCode.AGENT_CALL_ERROR.value,
                "error_msg": f"调用失败: {str(e)}"
            }
        except AgentNotFoundError:
            raise

    def _update_agent_task_metrics(self, instance: Dict[str, Any], success: bool) -> None:
        """更新Agent实例的任务指标（成功/失败次数、成功率）"""
        agent_type = instance.get("agent_type", "")
        instance_id = instance["instance_id"]
        try:
            # 初始化AgentState实例，更新指标
            agent_state = AgentState(
                agent_type=agent_type,
                instance_id=instance_id,
                redis_client=self.redis_client
            )
            agent_state.update_task_metrics(success=success)
            self.logger.debug(f"更新Agent实例{instance_id}指标成功: success={success}")
        except Exception as e:
            self.logger.warning(f"更新Agent实例{instance_id}指标失败: {str(e)}")
            # 指标更新失败不影响主流程，仅日志告警

    async def _retry_task_async(self, task_id: str, agent_type: str, task_params: Dict[str, Any]) -> Tuple[bool, Dict[str, Any]]:
        """异步任务重试逻辑（最多重试3次）"""
        task_manager = TaskStateManager(task_id=task_id, redis_client=self.redis_client)
        retry_count = task_manager.get_state()["retry_count"]
        max_retries = task_manager.get_state()["max_retries"]

        if retry_count >= max_retries:
            self.logger.error(f"任务{task_id}已达到最大重试次数({max_retries})，停止重试")
            task_manager.update_state(
                state=TaskState.FAILED.value,
                error_code=ErrorCode.TASK_EXECUTE_ERROR.value,
                error_msg=f"Max retries ({max_retries}) exceeded"
            )
            return False, {"success": False, "error_msg": "Max retries exceeded"}

        # 重试计数+1，更新任务状态为重试中
        task_manager.update_state(state=TaskState.RETRYING.value)
        self.logger.info(f"任务{task_id}开始第{retry_count + 1}次重试")

        try:
            # 重新选择Agent实例（避免复用失败的实例）
            instance = self._select_best_agent_instance(agent_type)
            # 重新调用
            result = await self._call_agent_instance_async(instance, task_params)
            if result["success"]:
                # 重试成功，更新任务状态
                task_manager.update_state(
                    state=TaskState.SUCCESS.value,
                    output_result=result["data"],
                    agent_instance_id=instance["instance_id"]
                )
                self._update_agent_task_metrics(instance, success=True)
                return True, result
            else:
                # 重试失败，继续重试或终止
                self._update_agent_task_metrics(instance, success=False)
                return await self._retry_task_async(task_id, agent_type, task_params)
        except Exception as e:
            self.logger.error(f"任务{task_id}第{retry_count + 1}次重试失败: {str(e)}")
            return await self._retry_task_async(task_id, agent_type, task_params)

    def _parse_task(self, task_input: Dict[str, Any]) -> List[Dict[str, Any]]:
        """解析任务（将复杂任务拆分为多个子任务，支持单/多Agent协作）"""
        task_type = task_input.get("task_type", "single")  # single/multi
        # 从嵌套的inputs对象中获取query
        input_data = task_input.get("inputs", {}).get("query", "")

        if task_type == "single":
            # 单Agent任务：直接指定Agent类型
            agent_type = task_input.get("agent_type", AgentType.TEXT_QA.value)
            return [{
                "sub_task_id": f"sub_{task_input['task_id']}_001",
                "agent_type": agent_type,
                "input": input_data,
                "task_id": task_input["task_id"],
                "session_id": task_input.get("inputs", {}).get("session_id", f"session_{task_input['task_id']}")
            }]
        elif task_type == "multi":
            # 多Agent协作任务：拆分为多个子任务（示例：工具调用→文本问答）
            return [
                {
                    "sub_task_id": f"sub_{task_input['task_id']}_001",
                    "agent_type": AgentType.TOOL_CALL.value,
                    "input": f"获取{input_data}相关信息",
                    "task_id": task_input["task_id"],
                    "session_id": task_input.get("inputs", {}).get("session_id", f"session_{task_input['task_id']}")
                },
                {
                    "sub_task_id": f"sub_{task_input['task_id']}_002",
                    "agent_type": AgentType.TEXT_QA.value,
                    "input": input_data,
                    "task_id": task_input["task_id"],
                    "session_id": task_input.get("inputs", {}).get("session_id", f"session_{task_input['task_id']}")
                }
            ]
        else:
            raise ParamError(
                message="Invalid task_type (only support 'single'/'multi')",
                context={"task_type": task_type, "error_code": ErrorCode.TASK_PARAM_ERROR.value}
            )

    async def execute_task_async(self, task_input: Dict[str, Any]) -> Dict[str, Any]:
        """异步执行任务（核心入口：解析→分发→执行→整合→重试）"""
        # 1. 校验任务输入参数
        required_fields = ["task_id", "inputs"]
        for field in required_fields:
            if field not in task_input:
                raise ParamError(
                    message=f"Task input missing required field: {field}",
                    context={"task_input": task_input, "error_code": ErrorCode.TASK_PARAM_ERROR.value}
                )
        # 确保inputs对象包含query字段
        if "query" not in task_input.get("inputs", {}):
            raise ParamError(
                message="Task input missing required field: query in inputs",
                context={"task_input": task_input, "error_code": ErrorCode.TASK_PARAM_ERROR.value}
            )
        task_id = task_input["task_id"]
        input_query = task_input["inputs"].get("query", "")
        self.logger.info(f"开始异步执行任务: task_id={task_id} | input={input_query}")

        # 2. 初始化任务状态（待执行）
        task_manager = TaskStateManager(task_id=task_id, redis_client=self.redis_client)
        task_manager.update_state(
            state=TaskState.PENDING.value,
            input_params=task_input
        )

        try:
            # 添加健康Agent实例信息日志
            with self.cache_lock:
                healthy_agent_info = {}
                for atype, instances in self.healthy_agents.items():
                    healthy_agent_info[atype] = len(instances)
            self.logger.info(f"当前健康Agent实例统计: {healthy_agent_info}")
            
            # 3. 解析任务（拆分子任务）
            sub_tasks = self._parse_task(task_input)
            self.logger.debug(
                f"任务解析完成，子任务数: {len(sub_tasks)} | sub_tasks={[t['sub_task_id'] for t in sub_tasks]}")

            # 4. 并发执行所有子任务（使用asyncio.gather实现真正的并行）
            sub_task_coroutines = []
            for sub_task in sub_tasks:
                sub_task_coroutines.append(self._execute_single_sub_task_async(sub_task, task_id))

            # 并发执行所有子任务
            sub_task_results = await asyncio.gather(*sub_task_coroutines, return_exceptions=True)

            # 处理执行结果
            valid_results = []
            for i, result in enumerate(sub_task_results):
                if isinstance(result, Exception):
                    self.logger.error(f"子任务{sub_tasks[i]['sub_task_id']}执行异常: {str(result)}")
                    # 可以选择重试或标记为失败
                    continue
                elif result and result.get("success"):
                    # 保持与_sub_task_results[0]["result"]期望的结构一致
                    valid_results.append({"result": result.get("data", {})})

            if not valid_results:
                raise AgentCallError(
                    message="All sub tasks failed",
                    context={"error_msg": "No valid results from any sub task"}
                )

            # 5. 整合所有子任务结果
            final_result = {
                "task_id": task_id,
                "status": "success",
                "sub_task_count": len(sub_tasks),
                "sub_task_results": valid_results,
                "final_answer": self._integrate_results(valid_results, task_input["inputs"].get("query", "")),
                "executed_at": int(datetime.now().timestamp())
            }

            # 记录任务完成和回顾相关日志
            final_answer = final_result['final_answer']
            if len(final_answer) > 50:
                log_answer = final_answer[:50] + '...'
            else:
                log_answer = final_answer
            self.logger.success(f"异步任务{task_id}执行完成并进行了结果回顾 | 最终结果: {log_answer}")
            return final_result
        except Exception as e:
            # 6. 任务执行失败：更新状态
            error_msg = str(e)
            self.logger.error(f"异步任务{task_id}执行失败: {error_msg}")
            task_manager.update_state(
                state=TaskState.FAILED.value,
                error_code=ErrorCode.TASK_EXECUTE_ERROR.value,
                error_msg=error_msg
            )
            return {
                "task_id": task_id,
                "status": "failed",
                "error_msg": error_msg,
                "executed_at": int(datetime.now().timestamp())
            }

    async def _execute_single_sub_task_async(self, sub_task: Dict[str, Any], task_id: str) -> Dict[str, Any]:
        """异步执行单个子任务"""
        sub_task_id = sub_task["sub_task_id"]
        agent_type = sub_task["agent_type"]
        self.logger.info(f"异步执行子任务: {sub_task_id} | agent_type={agent_type} | 查找最优Agent实例")

        # 更新子任务状态为执行中
        task_manager = TaskStateManager(task_id=task_id, redis_client=self.redis_client)
        task_manager.update_state(
            state=TaskState.RUNNING.value,
            agent_type=agent_type,
            input_params=sub_task
        )

        try:
            # 选择最优Agent实例
            instance = self._select_best_agent_instance(agent_type)
            if not instance:
                self.logger.error(f"找不到可用的{agent_type}类型Agent实例")
                raise AgentCallError(message=f"No available agent instance for type: {agent_type}")
            instance["agent_type"] = agent_type  # 补充Agent类型，用于更新指标
            self.logger.info(f"选中的Agent实例: instance_id={instance.get('instance_id')}, host={instance.get('host')}, port={instance.get('port')}")

            # 异步调用Agent实例执行子任务
            result = await self._call_agent_instance_async(instance, sub_task)

            if result["success"]:
                # 子任务成功：更新状态和Agent指标
                task_manager.update_state(
                    state=TaskState.SUCCESS.value,
                    agent_instance_id=instance["instance_id"],
                    output_result={"sub_task_id": sub_task_id, "result": result["data"]}
                )
                self._update_agent_task_metrics(instance, success=True)
                return result
            else:
                # 子任务失败：触发重试
                self.logger.error(f"子任务{sub_task_id}执行失败，触发重试: {result['error_msg']}")
                self._update_agent_task_metrics(instance, success=False)
                retry_success, retry_result = await self._retry_task_async(task_id, agent_type, sub_task)
                if retry_success:
                    return retry_result
                else:
                    raise AgentCallError(message=f"Sub task {sub_task_id} failed after retries",
                                         context={"error_msg": retry_result["error_msg"]})
        except Exception as e:
            self.logger.error(f"子任务{sub_task_id}执行异常: {str(e)}")
            raise

    def _integrate_results(self, sub_task_results: List[Dict[str, Any]], original_input: str) -> str:
        """整合子任务结果（包含任务回顾机制）"""
        if len(sub_task_results) == 1:
            # 单任务：直接返回结果
            result = sub_task_results[0]["result"]
        else:
            # 多任务：整合所有子任务结果
            integrated = f"针对查询「{original_input}」的综合结果：\n"
            for idx, result in enumerate(sub_task_results, 1):
                integrated += f"{idx}. {result['result']}\n"
            result = integrated.strip()
        
        # 执行任务回顾
        try:
            # 分析子任务结果
            analysis = task_reviewer.analyze_sub_task_results(sub_task_results)
            self.logger.debug(f"任务回顾分析结果: {analysis}")
            
            # 生成回顾摘要（同步方式）
            review_summary = task_reviewer.generate_review_summary_sync(
                original_query=original_input,
                analysis_result=analysis,
                sub_task_results=sub_task_results
            )
            
            # 评估结果质量
            quality_evaluation = task_reviewer.evaluate_result_quality(review_summary, original_input)
            
            # 记录回顾信息日志
            self.logger.info(f"任务回顾完成: quality_level={quality_evaluation.get('quality_level')}, "
                           f"redundancy_detected={len(analysis.get('redundant_pairs', []))}")
            
            # 如果回顾器认为结果有冗余或可以优化，返回优化后的摘要
            if quality_evaluation.get('quality_level') != '优秀' and review_summary:
                return review_summary
        except Exception as e:
            # 回顾失败不影响主流程
            self.logger.warning(f"任务回顾失败: {str(e)}")
        
        # 返回原始结果或整合结果
        return result


# ===================== 异步使用示例 =====================
async def main():
    """异步主函数示例"""
    import sys
    sys.path.append("../../")  # 添加项目根目录到Python路径

    # 初始化异步编排Agent
    async with AsyncOrchestratorAgent() as orchestrator:
        # 等待1秒，让订阅线程获取Agent实例缓存
        await asyncio.sleep(1)

        # 示例1：执行单Agent任务（文本问答）
        print("=== 执行异步单Agent任务（文本问答）===")
        single_task_input = {
            "task_id": "async_single_task_001",
            "task_type": "single",
            "agent_type": AgentType.TEXT_QA.value,
            "input": "什么是异步多Agent协作？"
        }
        single_result = await orchestrator.execute_task_async(single_task_input)
        print(json.dumps(single_result, ensure_ascii=False, indent=2))

        # 示例2：执行多Agent协作任务（并发执行）
        print("\n=== 执行异步多Agent协作任务（并发）===")
        multi_task_input = {
            "task_id": "async_multi_task_001",
            "task_type": "multi",
            "input": "请分析异步系统的优势"
        }
        multi_result = await orchestrator.execute_task_async(multi_task_input)
        print(json.dumps(multi_result, ensure_ascii=False, indent=2))


# 全局单例管理
_orchestrator_instance = None
_orchestrator_lock = asyncio.Lock()

async def get_async_orchestrator() -> AsyncOrchestratorAgent:
    """获取异步编排器实例（单例模式）"""
    global _orchestrator_instance
    async with _orchestrator_lock:
        if _orchestrator_instance is None:
            _orchestrator_instance = AsyncOrchestratorAgent()
    return _orchestrator_instance

if __name__ == "__main__":
    # 运行异步示例
    asyncio.run(main())